//
// Created by 神秘悟饭 on 2020/7/7.
//
#include "iostream"
#include <torch/script.h> // One-stop header.
#include <opencv2/opencv.hpp>
#include <iostream>
#include <opencv2/highgui.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>

using namespace std;
using namespace cv;
int classifier(string model_path,Mat image)
{
    int partType = 0;
    Mat image_transfomed;
    torch::jit::script::Module module = torch::jit::load(model_path);

    //输入图像
    cv::resize(image, image_transfomed, cv::Size(32, 32));
    // torch::from_blob无法读取uint16,故在此处归一化
    image_transfomed = image_transfomed / 65535;
    image_transfomed.convertTo(image_transfomed, CV_32F, 1, 0);

    // 图像转换为Tensor
    torch::Tensor tensor_image = torch::from_blob(image_transfomed.data, {image_transfomed.rows, image_transfomed.cols,1},torch::kFloat32);
    tensor_image = tensor_image.permute({2,0,1});
    tensor_image = tensor_image.unsqueeze(0);

    // 网络前向计算
    at::Tensor output = module.forward({tensor_image}).toTensor();

    auto max_result = output.max(1,true);
    auto max_index = std::get<1>(max_result).item<float>();

//    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

    // 返回分类标签 0未知，1颈椎，2脊椎，3胸椎，4髋部，5四肢
    // 原模型：0颈椎,1手，2踝，3髋，4膝
    if ((max_index == 1) || (max_index == 2) || (max_index == 4))
    {
        partType = 5;
    }
    else if (max_index == 3)
    {
        partType = 4;
    }
    else if (max_index == 0)
    {
        partType = 2;
    }
    return partType;

}
